# 对数据集resize 比letterbox效果好一点，没有阴影边

import os
import cv2
import glob
import numpy as np


def resize_image(img, new_shape=(640, 640)):
    resized_img = cv2.resize(img, new_shape, interpolation=cv2.INTER_LINEAR)
    return resized_img


def process_images(images_folder, labels_folder, output_folder):
    image_files = sorted(glob.glob(os.path.join(images_folder, '*.jpg'))) + sorted(
        glob.glob(os.path.join(images_folder, '*.bmp')))

    if not image_files:
        print("No image files found in the specified directory!")
        return

    os.makedirs(os.path.join(output_folder, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_folder, 'labels'), exist_ok=True)

    for image_file in image_files:
        base_name = os.path.splitext(os.path.basename(image_file))[0]
        label_file = os.path.join(labels_folder, base_name + '.txt')

        if not os.path.exists(label_file):
            print(f"Label file for {image_file} does not exist.")
            continue

        img = cv2.imread(image_file)
        if img is None:
            print(f"Failed to load image {image_file}")
            continue

        with open(label_file, 'r') as f:
            labels = [list(map(float, line.strip().split())) for line in f.readlines()]

        resized_img = resize_image(img)

        output_image_file = os.path.join(output_folder, 'images', os.path.basename(image_file))
        cv2.imwrite(output_image_file, resized_img)

        output_label_file = os.path.join(output_folder, 'labels', base_name + '.txt')
        with open(output_label_file, 'w') as f:
            for label in labels:
                class_id = int(label[0])
                coords = ' '.join(['{:.6f}'.format(coord) for coord in label[1:]])
                f.write(f"{class_id} {coords}\n")

        print(f"Processed {image_file}")


if __name__ == '__main__':
    images_folder = 'E:\Dataset/adddata\images/test/dark_bright'
    labels_folder = 'E:\Dataset/adddata\labels/test/dark_bright'
    output_folder = 'E:\Dataset/adddata/test/dark_bright'

    process_images(images_folder, labels_folder, output_folder)
